// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <algorithm>

#include <pollux/vector/simple_vector.h>

namespace kumo::pollux {
    template<typename T, typename ShortType>
    BufferPtr build_biased_buffer(
        const std::vector<std::optional<T> > &values,
        T bias,
        pollux::memory::MemoryPool *pool) {
        BufferPtr buffer = AlignedBuffer::allocate<ShortType>(values.size(), pool);
        ShortType *data = buffer->asMutable<ShortType>();
        buffer->setSize(sizeof(ShortType) * values.size());

        for (vector_size_t i = 0; i < values.size(); ++i) {
            if (values[i] != std::nullopt) {
                data[i] = static_cast<ShortType>(*values[i] - bias);
            }
        }
        return buffer;
    }

    template<typename T>
    BiasVectorPtr<EvalType<T> > VectorMaker::bias_vector(
        const std::vector<std::optional<T> > &data) {
        POLLUX_CHECK_GT(data.size(), 1);
        using TEvalType = EvalType<T>;

        if constexpr (admitsBias<TEvalType>()) {
            auto stats = genVectorMakerStats(data);
            POLLUX_CHECK(stats.min != std::nullopt);
            POLLUX_CHECK(stats.max != std::nullopt);
            TEvalType min = *stats.min;
            TEvalType max = *stats.max;

            // This ensures the math conversions when getting a delta on signed
            // values at opposite ends of the int64 range do not overflow a
            // temporary signed value.
            uint64_t delta = max < 0
                                 ? max - min
                                 : min < 0
                                       ? static_cast<uint64_t>(max) + static_cast<uint64_t>(std::abs(min))
                                       : max - min;

            POLLUX_CHECK(deltaAllowsBias<TEvalType>(delta));

            // Check BiasVector.h for explanation of this calculation.
            TEvalType bias = min + static_cast<TEvalType>(std::ceil(delta / 2.0));

            BufferPtr buffer;
            TypeKind valueType;

            if (delta <= std::numeric_limits<uint8_t>::max()) {
                buffer = build_biased_buffer<TEvalType, int8_t>(data, bias, pool_);
                valueType = TypeKind::TINYINT;
            } else if (delta <= std::numeric_limits<uint16_t>::max()) {
                buffer = build_biased_buffer<TEvalType, int16_t>(data, bias, pool_);
                valueType = TypeKind::SMALLINT;
            } else {
                buffer = build_biased_buffer<TEvalType, int32_t>(data, bias, pool_);
                valueType = TypeKind::INTEGER;
            }

            auto bias_vector = std::make_shared<BiasVector<TEvalType> >(
                pool_,
                nullptr /*nulls*/,
                data.size(),
                valueType,
                buffer,
                bias,
                stats.asSimpleVectorStats(),
                stats.distinctCount(),
                stats.nullCount,
                stats.isSorted);

            for (vector_size_t i = 0; i < data.size(); i++) {
                if (data[i] == std::nullopt) {
                    bias_vector->set_null(i, true);
                }
            }
            return bias_vector;
        } else {
            POLLUX_UNSUPPORTED("Invalid type for biasing");
        }
    }

    template<typename T, typename TEvalType>
    void sequenceEncode(
        const std::vector<std::optional<T> > &input,
        std::vector<std::optional<TEvalType> > &encodedVals,
        std::vector<SequenceLength> &encodedLengths) {
        POLLUX_CHECK_GT(input.size(), 0, "need at least one element to encode.");
        auto currentValue = input.front();
        SequenceLength currentRun = 1;

        for (vector_size_t i = 1; i < input.size(); i++) {
            if (input[i] == currentValue) {
                ++currentRun;
            } else {
                encodedVals.emplace_back(currentValue);
                encodedLengths.emplace_back(currentRun);
                currentValue = input[i];
                currentRun = 1;
            }
        }
        encodedVals.emplace_back(currentValue);
        encodedLengths.emplace_back(currentRun);
    }

    template<typename T>
    SequenceVectorPtr<EvalType<T> > VectorMaker::sequence_vector(
        const std::vector<std::optional<T> > &data) {
        using TEvalType = EvalType<T>;

        std::vector<std::optional<TEvalType> > sequenceVals;
        std::vector<SequenceLength> sequenceLengths;
        sequenceEncode(data, sequenceVals, sequenceLengths);

        auto stats = genVectorMakerStats(data);
        return std::make_unique<SequenceVector<TEvalType> >(
            pool_,
            data.size(),
            flat_vector_nullable(sequenceVals),
            copyToBuffer(sequenceLengths, pool_),
            stats.asSimpleVectorStats(),
            stats.distinctCount(),
            stats.nullCount,
            stats.isSorted);
    }

    template<typename T>
    ConstantVectorPtr<EvalType<T> > VectorMaker::constant_vector(
        const std::vector<std::optional<T> > &data) {
        POLLUX_CHECK_GT(data.size(), 0);
        using TEvalType = EvalType<T>;

        auto stats = genVectorMakerStats(data);
        vector_size_t distinctCount = stats.distinctCount();
        vector_size_t nullCount = stats.nullCount;

        POLLUX_CHECK(
            (distinctCount == 1 && nullCount == 0) || distinctCount == 0,
            "Attempting to build a constant vector with invalid entries");

        return std::make_unique<ConstantVector<TEvalType> >(
            pool_,
            data.size(),
            nullCount > 0,
            CppToType<TEvalType>::create(),
            (nullCount > 0) ? TEvalType() : melon::copy(*data.front()),
            stats.asSimpleVectorStats());
    }

    template<typename T>
    DictionaryVectorPtr<EvalType<T> > VectorMaker::dictionary_vector(
        const std::vector<std::optional<T> > &data) {
        using TEvalType = EvalType<T>;

        // Encodes the data saving distinct values on `distinctValues` and their
        // respective indices on `indices`.
        std::vector<TEvalType> distinctValues;
        melon::F14FastMap<TEvalType, int32_t> indexMap;

        BufferPtr indices = AlignedBuffer::allocate<int32_t>(data.size(), pool_);
        auto rawIndices = indices->asMutable<int32_t>();

        BufferPtr nulls =
                AlignedBuffer::allocate<bool>(data.size(), pool_, bits::kNotNull);
        auto rawNulls = nulls->asMutable<uint64_t>();

        vector_size_t nullCount = 0;

        for (auto i = 0; i < data.size(); ++i) {
            auto val = data[i];
            if (val == std::nullopt) {
                ++nullCount;
                bits::set_null(rawNulls, i, true);
            } else {
                const auto &[it, inserted] = indexMap.emplace(*val, indexMap.size());
                if (inserted) {
                    distinctValues.push_back(*val);
                }
                *rawIndices = it->second;
            }
            ++rawIndices;
        }

        auto values = flat_vector(distinctValues);
        auto stats = genVectorMakerStats(data);
        auto dictionary_vector = std::make_unique<DictionaryVector<TEvalType> >(
            pool_,
            nullCount ? nulls : nullptr,
            data.size(),
            std::move(values),
            std::move(indices),
            stats.asSimpleVectorStats(),
            indexMap.size(),
            nullCount,
            stats.isSorted);

        return dictionary_vector;
    }

    template<typename T>
    FlatVectorPtr<EvalType<T> > VectorMaker::flat_vector_nullable(
        const std::vector<std::optional<T> > &data,
        const TypePtr &type) {
        using TEvalType = EvalType<T>;
        BufferPtr dataBuffer = AlignedBuffer::allocate<TEvalType>(data.size(), pool_);

        auto stats = genVectorMakerStats(data);
        auto flat_vector = std::make_shared<FlatVector<TEvalType> >(
            pool_,
            type,
            BufferPtr(nullptr),
            data.size(),
            std::move(dataBuffer),
            std::vector<BufferPtr>(),
            stats.asSimpleVectorStats(),
            stats.distinctCount(),
            stats.nullCount,
            stats.isSorted);

        for (vector_size_t i = 0; i < data.size(); i++) {
            if (data[i] != std::nullopt) {
                flat_vector->set(i, TEvalType(*data[i]));
                flat_vector->set_null(i, false);
            } else {
                flat_vector->set(i, TEvalType());
                flat_vector->set_null(i, true);
            }
        }
        return flat_vector;
    }

    template<typename T>
    FlatVectorPtr<EvalType<T> > VectorMaker::flat_vector(
        const std::vector<T> &data,
        const TypePtr &type) {
        using TEvalType = EvalType<T>;
        BufferPtr dataBuffer = AlignedBuffer::allocate<TEvalType>(data.size(), pool_);

        auto stats = genVectorMakerStats(data);
        auto flat_vector = std::make_shared<FlatVector<TEvalType> >(
            pool_,
            type,
            BufferPtr(nullptr),
            data.size(),
            std::move(dataBuffer),
            std::vector<BufferPtr>(),
            stats.asSimpleVectorStats(),
            stats.distinctCount(),
            stats.nullCount,
            stats.isSorted);

        for (vector_size_t i = 0; i < data.size(); i++) {
            flat_vector->set(i, TEvalType(data[i]));
        }
        return flat_vector;
    }
} // namespace kumo::pollux
